import os
import torch
import sys
from pathlib import Path
sys.path.append(str(Path(__file__).parent.parent))
from torch import nn
from torchvision.models import resnet50
os.environ["CUDA_VISIBLE_DEVICES"]="0"
from model_effcientNet import dataset_eight_slice

class Backbone(nn.Module):
    def __init__(self):
        super().__init__()
        base_model = resnet50(pretrained=False)
        encoder_layers = list(base_model.children())
        self.backbone = nn.Sequential(*encoder_layers[:8])

    def forward(self, x):
        return self.backbone(x)

class ModalitySpecificNet(nn.Module):
    def __init__(self):
        super().__init__()
        # 每个模态独立的主干网络
        self.forme_backbone = Backbone() # 期待B，3，H，W的输入
        self.forme_backbone.load_state_dict(torch.load("./pretrained_model/RadImageNet_pytorch/ResNet50.pt"))
        self.mid_backbone = Backbone()
        self.mid_backbone.load_state_dict(torch.load("./pretrained_model/RadImageNet_pytorch/ResNet50.pt"))
        self.back_backbone = Backbone()
        self.back_backbone.load_state_dict(torch.load("./pretrained_model/RadImageNet_pytorch/ResNet50.pt"))
        # 需要[b,1024,14,14]的数据

    def forward(self, imgs):
        forme_img = imgs[:,0,:,:]
        if forme_img.size(1) == 1:
            forme_img = forme_img.repeat(1,3,1,1)
        mid_img = imgs[:,1,:,:]
        if mid_img.size(1) == 1:
            mid_img = mid_img.repeat(1,3,1,1)
        back_img = imgs[:,2,:,:]
        if back_img.size(1) == 1:
            back_img = back_img.repeat(1,3,1,1)

        forme_feature = self.forme_backbone(forme_img)
        mid_feature = self.mid_backbone(mid_img)
        back_feature = self.back_backbone(back_img)
        # 特征拼接，有无其他方法，比如利用ROI做拼接
        feature = torch.cat([forme_feature, mid_feature, back_feature], dim=1)
        return feature
if __name__ == "__main__":
    model = ModalitySpecificNet()
    data = torch.randn(4, 3, 224, 224)
    print(model(data).shape)
    